import os
import shutil
from pathlib import Path
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--pcd', default='/media/zw_0725/view/cloud', type=str, metavar='PATH',
                    help='path to pcd dir')
parser.add_argument('--img', default='/media/zw_0725/view/image', type=str, metavar='PATH',
                    help='path to pcd dir')
parser.add_argument('--save', default='/media/zw_0725/view/image_sync', type=str, metavar='PATH',
                    help='path to pcd dir')

args = parser.parse_args()

def find_closest_timestamp(target, timestamps):
    """
    Find the closest timestamp to the target timestamp from a list of timestamps.
    """
    return min(timestamps, key=lambda x: abs(x - target))

def main(point_cloud_dir, image_dir, output_image_dir):
    # Convert to Path objects for easier manipulation
    point_cloud_dir = Path(point_cloud_dir)
    image_dir = Path(image_dir)
    output_image_dir = Path(output_image_dir)
    
    # Ensure output directory exists
    output_image_dir.mkdir(parents=True, exist_ok=True)
    
    # Extract timestamps from file names
    point_cloud_files = list(point_cloud_dir.glob("cloud_*.pcd"))
    image_files = list(image_dir.glob("*.png"))
    
    point_cloud_timestamps = [int(f.stem.split('_')[1]) for f in point_cloud_files]
    image_timestamps = [int(f.stem) for f in image_files]
    
    image_files_dict = {int(f.stem): f for f in image_files}
    
    # Process each point cloud file
    for pc_file in point_cloud_files:
        pc_timestamp = int(pc_file.stem.split('_')[1])
        
        # Find the closest image timestamp
        closest_image_timestamp = find_closest_timestamp(pc_timestamp, image_timestamps)
        
        # Get the corresponding image file
        closest_image_file = image_files_dict[closest_image_timestamp]
        
        # Construct the output file name
        output_file_name = f"cloud_{pc_timestamp}.png"
        output_file_path = output_image_dir / output_file_name
        
        # Copy the image file to the new location with the new name
        shutil.copy(closest_image_file, output_file_path)
        print(f"Copied {closest_image_file} to {output_file_path}")

if __name__ == "__main__":
    point_cloud_dir = args.pcd
    image_dir = args.img
    output_image_dir = args.save
    
    main(point_cloud_dir, image_dir, output_image_dir)

